// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/vector_function.h>
#include <pollux/type/floating_point_util.h>
#include <pollux/vector/decoded_vector.h>

namespace kumo::pollux::functions {
    namespace {
        template<typename T>
        inline bool isPrimitiveEqual(const T &lhs, const T &rhs) {
            if constexpr (std::is_floating_point_v<T>) {
                return util::floating_point::NaNAwareEquals<T>{}(lhs, rhs);
            } else {
                return lhs == rhs;
            }
        }

        template<TypeKind kind>
        void applyTyped(
            const SelectivityVector &rows,
            DecodedVector &arrayDecoded,
            DecodedVector &elementsDecoded,
            DecodedVector &searchDecoded,
            exec::EvalCtx & /*context*/,
            FlatVector<bool> &flatResult,
            bool /*throwOnNestedNull*/) {
            using T = typename TypeTraits<kind>::NativeType;

            auto baseArray = arrayDecoded.base()->as<ArrayVector>();
            auto rawSizes = baseArray->rawSizes();
            auto rawOffsets = baseArray->rawOffsets();
            auto indices = arrayDecoded.indices();

            constexpr bool isBoolType = std::is_same_v<bool, T>;

            if (!isBoolType && elementsDecoded.isIdentityMapping() &&
                !elementsDecoded.may_have_nulls() && searchDecoded.isConstantMapping()) {
                auto rawElements = elementsDecoded.data<T>();
                auto search = searchDecoded.value_at<T>(0);

                rows.applyToSelected([&](auto row) {
                    auto size = rawSizes[indices[row]];
                    auto offset = rawOffsets[indices[row]];

                    for (auto i = 0; i < size; i++) {
                        if (isPrimitiveEqual<T>(rawElements[offset + i], search)) {
                            flatResult.set(row, true);
                            return;
                        }
                    }
                    flatResult.set(row, false);
                });
            } else {
                rows.applyToSelected([&](auto row) {
                    auto size = rawSizes[indices[row]];
                    auto offset = rawOffsets[indices[row]];

                    auto search = searchDecoded.value_at<T>(row);

                    bool foundNull = false;

                    for (auto i = 0; i < size; i++) {
                        if (elementsDecoded.is_null_at(offset + i)) {
                            foundNull = true;
                        } else if (isPrimitiveEqual<T>(
                            elementsDecoded.value_at<T>(offset + i), search)) {
                            flatResult.set(row, true);
                            return;
                        }
                    }

                    if (foundNull) {
                        flatResult.set_null(row, true);
                    } else {
                        flatResult.set(row, false);
                    }
                });
            }
        }

        void applyComplexType(
            const SelectivityVector &rows,
            DecodedVector &arrayDecoded,
            DecodedVector &elementsDecoded,
            DecodedVector &searchDecoded,
            exec::EvalCtx &context,
            FlatVector<bool> &flatResult,
            bool throwOnNestedNull) {
            auto baseArray = arrayDecoded.base()->as<ArrayVector>();
            auto rawSizes = baseArray->rawSizes();
            auto rawOffsets = baseArray->rawOffsets();
            auto indices = arrayDecoded.indices();

            auto elementsBase = elementsDecoded.base();
            auto elementIndices = elementsDecoded.indices();
            auto searchBase = searchDecoded.base();
            auto searchIndices = searchDecoded.indices();

            const auto nullHandlingMode = throwOnNestedNull
                                              ? CompareFlags::NullHandlingMode::kNullAsIndeterminate
                                              : CompareFlags::NullHandlingMode::kNullAsValue;

            context.applyToSelectedNoThrow(rows, [&](auto row) {
                auto size = rawSizes[indices[row]];
                auto offset = rawOffsets[indices[row]];
                bool foundNull = false;
                auto searchIndex = searchIndices[row];

                for (auto i = 0; i < size; i++) {
                    if (elementsDecoded.is_null_at(offset + i)) {
                        foundNull = true;
                    } else {
                        std::optional<bool> result = elementsBase->equal_value_at(
                            searchBase,
                            elementIndices[offset + i],
                            searchIndex,
                            nullHandlingMode);
                        POLLUX_USER_CHECK(
                            result.has_value(),
                            "contains does not support arrays with elements that contain null");
                        if (result.value()) {
                            flatResult.set(row, true);
                            return;
                        }
                    }
                }

                if (foundNull) {
                    flatResult.set_null(row, true);
                } else {
                    flatResult.set(row, false);
                }
            });
        }

        template<>
        void applyTyped<TypeKind::ARRAY>(
            const SelectivityVector &rows,
            DecodedVector &arrayDecoded,
            DecodedVector &elementsDecoded,
            DecodedVector &searchDecoded,
            exec::EvalCtx &context,
            FlatVector<bool> &flatResult,
            bool throwOnNestedNull) {
            applyComplexType(
                rows,
                arrayDecoded,
                elementsDecoded,
                searchDecoded,
                context,
                flatResult,
                throwOnNestedNull);
        }

        template<>
        void applyTyped<TypeKind::MAP>(
            const SelectivityVector &rows,
            DecodedVector &arrayDecoded,
            DecodedVector &elementsDecoded,
            DecodedVector &searchDecoded,
            exec::EvalCtx &context,
            FlatVector<bool> &flatResult,
            bool throwOnNestedNull) {
            applyComplexType(
                rows,
                arrayDecoded,
                elementsDecoded,
                searchDecoded,
                context,
                flatResult,
                throwOnNestedNull);
        }

        template<>
        void applyTyped<TypeKind::ROW>(
            const SelectivityVector &rows,
            DecodedVector &arrayDecoded,
            DecodedVector &elementsDecoded,
            DecodedVector &searchDecoded,
            exec::EvalCtx &context,
            FlatVector<bool> &flatResult,
            bool throwOnNestedNull) {
            applyComplexType(
                rows,
                arrayDecoded,
                elementsDecoded,
                searchDecoded,
                context,
                flatResult,
                throwOnNestedNull);
        }

        class ArrayContainsFunction : public exec::VectorFunction {
        public:
            explicit ArrayContainsFunction(bool throwOnNestedNull)
                : throwOnNestedNull_(throwOnNestedNull) {
            }

            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr & /* outputType */,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                POLLUX_CHECK_EQ(args.size(), 2);
                const auto &array_vector = args[0];
                const auto &searchVector = args[1];

                POLLUX_CHECK(array_vector->type()->is_array());
                POLLUX_CHECK(array_vector->type()->as_array().elementType()->kindEquals(
                    searchVector->type()));

                context.ensure_writable(rows, BOOLEAN(), result);
                auto flatResult = result->as_flat_vector<bool>();

                exec::LocalDecodedVector arrayHolder(context, *array_vector, rows);
                auto elements = arrayHolder.get()->base()->as<ArrayVector>()->elements();

                exec::LocalSelectivityVector nestedRows(context, elements->size());
                nestedRows.get()->setAll();

                exec::LocalDecodedVector elementsHolder(
                    context, *elements, *nestedRows.get());

                exec::LocalDecodedVector searchHolder(context, *searchVector, rows);

                if (searchVector->type()->providesCustomComparison()) {
                    // We use applyComplexType for types that provide custom comparison
                    // operators because the main difference between applyComplexType and
                    // applyTyped is that applyComplexType calls the Vector's equal_value_at
                    // method, which calls the Types custom comparison operator internally.
                    applyComplexType(
                        rows,
                        *arrayHolder.get(),
                        *elementsHolder.get(),
                        *searchHolder.get(),
                        context,
                        *flatResult,
                        throwOnNestedNull_);
                } else {
                    POLLUX_DYNAMIC_TYPE_DISPATCH(
                        applyTyped,
                        searchVector->type_kind(),
                        rows,
                        *arrayHolder.get(),
                        *elementsHolder.get(),
                        *searchHolder.get(),
                        context,
                        *flatResult,
                        throwOnNestedNull_);
                }
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                // array(T), T -> boolean
                return {
                    exec::FunctionSignatureBuilder()
                    .typeVariable("T")
                    .returnType("boolean")
                    .argumentType("array(T)")
                    .argumentType("T")
                    .build()
                };
            }

        private:
            bool throwOnNestedNull_;
        };
    } // namespace

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_array_contains,
        ArrayContainsFunction::signatures(),
        std::make_unique<ArrayContainsFunction>(true));

    // Internal function only used for testing. This function allows the array to
    // have null elements and considers null as a value, i.e., null == null.
    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_$internal$contains,
        ArrayContainsFunction::signatures(),
        std::make_unique<ArrayContainsFunction>(false));
} // namespace kumo::pollux::functions
